package cn.zaink;

import cn.hutool.core.lang.Assert;
import cn.hutool.core.util.StrUtil;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.jmeter.config.Arguments;
import org.apache.jmeter.protocol.java.sampler.AbstractJavaSamplerClient;
import org.apache.jmeter.protocol.java.sampler.JavaSamplerContext;

import java.util.concurrent.CountDownLatch;

import static cn.zaink.AbstractZkClient.Arg.*;

/**
 * @author zenghy
 **/
@Slf4j
public abstract class AbstractZkClient extends AbstractJavaSamplerClient {

    public interface Arg {
        String SERVER = "Server";
        String COUNTER = "Counter";
        String R_COUNTER = "RCounter";
        String W_COUNTER = "WCounter";
        String RW_COUNTER = "RwCounter";

        String NODE_PREFIX = "NodePrefix";
        String NODE_CHILD_PREFIX = "NodeChildPrefix";
        String TREE_LEVEL = "TreeLevel";
    }

    CuratorFramework client;
    CountDownLatch latch;

    @Override
    public Arguments getDefaultParameters() {
        Arguments arguments = new Arguments();
        arguments.addArgument(SERVER, "host:port");
        arguments.addArgument(COUNTER, "${__counter(FALSE,counter)}");
        arguments.addArgument(R_COUNTER, "${__counter(TRUE,RCounter)}");
        arguments.addArgument(W_COUNTER, "${__counter(TRUE,WCounter)}");
        arguments.addArgument(RW_COUNTER, "${__counter(TRUE,RwCounter)}");
        arguments.addArgument(NODE_PREFIX, "/node");
        arguments.addArgument(NODE_CHILD_PREFIX, "/childNode");
        arguments.addArgument(TREE_LEVEL, "1");
        return arguments;
    }

    @Override
    public void setupTest(JavaSamplerContext context) {
        // 这个方法只会运行一次
        String server = context.getParameter(SERVER);
        Assert.isTrue(StrUtil.isNotBlank(server), "zk服务器地址不能为空");
        latch = new CountDownLatch(1);
        ExponentialBackoffRetry retryPolicy = new ExponentialBackoffRetry(1000, 3);
        client = CuratorFrameworkFactory.newClient(server, retryPolicy);
        log.warn("连接到zk服务器： {}", server);
        client.start();
        Assert.isTrue(client.getState().equals(CuratorFrameworkState.STARTED), "连接zk服务器失败");
    }

    @SneakyThrows
    @Override
    public void teardownTest(JavaSamplerContext javaSamplerContext) {
        if (null != client) {
            CuratorFrameworkState state = client.getState();
            if (state.equals(CuratorFrameworkState.STARTED)) {
                latch.await();
                log.warn("断开zk连接");
                client.close();
            }
        }
    }
}
